from langchain.chains.router import MultiPromptChain
from langchain_community.llms import Tongyi
from langchain_core.runnables.base import RunnableSequence
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain.prompts import PromptTemplate
from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE

from dotenv import load_dotenv, find_dotenv

load_dotenv(find_dotenv())

physics_prompt_tpl = PromptTemplate.from_template(
    '你是个优秀的物理学家，你很擅长用简明易懂的方式回答有关物理学的问题。'
    '请使用中文帮我解答下列问题：\n{input}'
)
math_prompt_tpl = PromptTemplate.from_template(
    '你是个很好的数学家。你很擅长回答数学问题。'
    '请使用中文帮我解答下列问题：\n{input}'
)

# 创建模板信息列表
prompt_infos = [
    {
        'name': 'physics',
        'description': '用于解答物理相关问题',
        'prompt_template': physics_prompt_tpl,
    },
    {
        'name': 'math',
        'description': '用于解答数学相关问题',
        'prompt_template': math_prompt_tpl,
    },
]
llm = Tongyi(temperature=0.1)

# 生成键为模板名称、值为Chain的字典
destination_chains = {}
for p_info in prompt_infos:
    name = p_info['name']
    prompt = p_info['prompt_template']
    chain = RunnableSequence(prompt | llm)
    destination_chains[name] = chain

# 将模板名称和模板描述通过MULTI_PROMPT_ROUTER_TEMPLATE生成模板
destinations = [f'{p["name"]}: {p["description"]}' for p in prompt_infos]
destinations_str = "\n".join(destinations)

router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(destinations=destinations_str)
router_prompt = PromptTemplate(
    template=router_template,
    input_variables=['input'],
    output_parser=RouterOutputParser(),
)
router_chain = LLMRouterChain.from_llm(llm, router_prompt)

# 这里创建了一个default_chain
# 为了防止提的问题类型并没有包含在prompt_infos中
default_chain = RunnableWithMessageHistory(llm=llm, output_key='text')
chain = MultiPromptChain(
    router_chain=router_chain,
    destination_chains=destination_chains,
    default_chain=default_chain,
    verbose=True,
)

print(chain.invoke("什么是牛顿第一定律?"))
print(chain.invoke("今晚有鸡兔同笼，上有35个头，下有94只脚。鸡兔各有几只？"))